"""
    Computational Economics
    2.8: Schelling's Segregation Model
    http://johnstachurski.net/lectures/schelling.html

    DEFINITIONS

    QUOTES
    "[Y]ou're standing at the edge of a cliff, chained by the ankle to someone
    else. You'll be released, and one of you will get a large prize, as soon as
    the other gives in. How do you persuade the other guy to give in, when the
    only method at your disposal -- threatening to push him off the cliff --
    would doom you both?"

    "Answer: You start dancing, closer and closer to the edge. That way, you
    don't have to convince him that you would do something totally irrational:
    plunge him and yourself off the cliff. You just have to convince him that
    you are prepared to take a higher risk than he is of accidentally falling
    off the cliff. If you can do that, you win."

    REFERENCES
    http://en.wikipedia.org/wiki/Thomas_Schelling

    OUTLINE OF MODEL
    Suppose we have two types of people

        * Orange people and green people

    The live mixed together on the unit square

        * Locations of the form (x, y), where 0 < x, y < 1

    Initially they are integrated

        * Starting locations IID bivariate uniform on unit square

    Each agent is now given the chance to stay or move

        * Stay if they are "happy," move if they are "unhappy"
        * Happy if half or more of 10 nearest neighbors are of the same type
        * Nearest is in terms of Euclidean distance

    If agent is unhappy

        * Draw random location
              o If happy at new location, move there
              o Else, repeat

    We cycle through the agents, giving offers to move

    Continuing until no-one wishes to move

    MODELING
    Agents are modeled as objects

        * Data:
              o type and location
        * Methods:
              o Determine whether happy or not given locations of other agents
              o If not happy, move
                    + find a new location where happy

    Pseudocode for the main loop

    while agents are still moving:
        for agent in agents:
            give agent the opportunity to move

    Use 200 agents of each type

    Plot the before and after locations

"""

import random, pylab
import unittest
from math import sqrt
import time
from os.path import (dirname, join)

class Agent:

    type = ''
    start_location = [0,0]
    location = [0, 0]
    locations = []

    def __init__(self, type, location=None):
        self.type = type
        if not location:
            self.start_location = self.get_random_location()
        else:
            self.start_location = location
        self.location = self.start_location
        self.num_relocations = 0

    def __str__(self):
        return '<Agent id=%s type=%s location=%s relocations=%s>' % (
            id(self), self.type, tuple(self.location), len(self.locations) )

    def is_happy(self, other_agents):
        num_neighbors = 10
        neighbors = self.find_closest_neighbors(other_agents, num_neighbors)
        same_type = [1 for n, d in neighbors if n.type == self.type]
        return sum(same_type) >= (num_neighbors / 2)

    def relocate(self, all_agents):
        while not self.is_happy(all_agents):
            self.location = self.get_random_location()
            self.locations.append(self.location)

    def find_closest_neighbors(self, neighbors, num):
        neighbor_distances = []
        for n in neighbors:
            if n is self:
                continue
            d = self.measure_distance(n.location)
            neighbor_distances.append((n, d))
        closest_neighbors = sorted(neighbor_distances, key=lambda tup: tup[1], reverse=False)
        return closest_neighbors[:num]

    def measure_distance(self, location):
        x1, y1 = location
        x0, y0 = self.location
        return sqrt((x1-x0)**2 + (y1-y0)**2)

    def get_random_location(self):
        return (random.uniform(0, 1), random.uniform(0, 1))


def draw_map(agents, title):
    pylab.title('Schelling Model: %s' % (title))
    pylab.grid(True)

    for agent in agents:
        format = '%so' % (agent.type[0].lower())
        x, y = agent.location
        pylab.plot([x], [y], format)

    fig_title = 'fig-%s.png' % (title)
    path = join(dirname(__file__), fig_title)
    pylab.savefig(path)
    pylab.clf()


class AgentTest(unittest.TestCase):
    def setUp(self):
        self.agent = Agent('test', [0,0])

    def tearDown(self):
        self.agent = None

    def testRelocation(self):
        is_grouchy = lambda l: random.randint(1, 100) < 3
        self.agent.is_happy = is_grouchy
        location1 = self.agent.location
        self.agent.relocate([])
        print self.agent
        self.assertNotEqual(self.agent.location, location1)

    def testMeasureDistance(self):
        cases = [
            ((1,1), sqrt(2)),
            ((2,2), sqrt(8))
        ]
        for location, expect in cases:
            self.assertEqual(self.agent.measure_distance(location), expect)

    def testIdentity(self):
        other = Agent('test', [0,0])
        self.assertTrue(self.agent is self.agent)
        self.assertFalse(self.agent is other)

    def testAgentConstructor(self):
        self.assertEqual(self.agent.type, 'test')
        self.assertEqual(self.agent.location, [0,0])


def agents_relocate(agents):
    t = time.time()
    random.shuffle(agents)
    num_relocations = 0
    for agent in agents:
        if not agent.is_happy(agents):
            agent.relocate(agents)
            num_relocations += 1

    tr = time.time() - t
    print '[%.2fs] %s of %s agents relocated this round' % (tr,
        num_relocations, len(agents))
    return num_relocations > 0


def main():
    num_agents = 200
    red_agents = []
    blue_agents = []
    stage = 0

    # create red agents
    for n in range(num_agents):
        red_agents.append(Agent('red'))

    # create blue agents
    for n in range(num_agents):
        blue_agents.append(Agent('blue'))

    # draw starting map
    agents = red_agents + blue_agents
    random.shuffle(agents)
    draw_map(agents, 'initial')

    # move
    while agents_relocate(agents):
        stage += 1
        draw_map(agents, 'stage-%s' % (stage))
        print random.choice(agents)

    # draw ending map
    random.shuffle(agents)
    draw_map(agents, 'final')

#
# MAIN
#
if __name__ == "__main__":
    suite = unittest.TestLoader().loadTestsFromTestCase(AgentTest)
    result = unittest.TextTestRunner(verbosity=2).run(suite)
    if result.errors:
        raise Exception('error in tests')
    main()
    print '%s: ok' % (__file__)
